import numpy as np
import cv2 as cv
import glob

# 迭代终止条件
criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
# 设置目标点：(0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
objp = np.zeros((6*7,3), np.float32)
objp[:,:2] = np.mgrid[0:7,0:6].T.reshape(-1,2)
objp = objp
# 对象点和图像点列表
objpoints = [] # 真实世界空间的3d点
imgpoints = [] # 图像平面的2d点
images = glob.glob('./chess_board_images/*.jpg')
for fname in images:
    img = cv.imread(fname)
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    # 查找棋盘点
    ret, corners = cv.findChessboardCorners(gray, (7,6), None)
    # 如果找到棋盘点，加入对象点列表，图像列表
    if ret == True:
        objpoints.append(objp)
        corners2 = cv.cornerSubPix(gray,corners, (11,11), (-1,-1), criteria)
        imgpoints.append(corners2)
        # 绘制角点
        cv.drawChessboardCorners(img, (7,6), corners2, ret)
        cv.imshow('img', img)
        cv.waitKey(500)
cv.destroyAllWindows()


ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(objpoints,
                                                  imgpoints,
                                                  gray.shape[::-1],
                                                  None, None)

print("camera matrix :\n",mtx)
print("distortion coefficients :\n",dist)

np.savez("./data/camera_param.npz",mtx=mtx,dist=dist,rvecs=rvecs,tvecs=tvecs)